package org.cybzacg.encryption.utils;

import org.cybzacg.encryption.enums.EncryptionAlgorithm;
import org.cybzacg.encryption.enums.EncryptionType;
import org.cybzacg.encryption.model.EncryptionResult;

import javax.crypto.*;
import javax.crypto.spec.*;
import java.security.*;
import java.security.spec.AlgorithmParameterSpec;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Arrays;
import java.util.Base64;

/**
 * 加密工具类
 * 提供各种加密算法的通用加密/解密功能
 * 
 * @author encryption team
 * @version 1.0
 */
public class EncryptionUtils {
    
    private static final String AES_TRANSFORMATION = "AES/CBC/PKCS5Padding";
    private static final String DES_TRANSFORMATION = "DES/CBC/PKCS5Padding";
    private static final String TRIPLE_DES_TRANSFORMATION = "DESede/CBC/PKCS5Padding";
    private static final String BLOWFISH_TRANSFORMATION = "Blowfish/CBC/PKCS5Padding";
    private static final String RSA_TRANSFORMATION = "RSA/ECB/PKCS1Padding";
    private static final String RSA_OAEP_TRANSFORMATION = "RSA/ECB/OAEPWithSHA-256AndMGF1Padding";
    
    /**
     * AES加密
     * 
     * @param data 原始数据
     * @param key 密钥
     * @param iv 初始化向量
     * @return 加密结果
     * @throws GeneralSecurityException 加密异常
     */
    public static EncryptionResult encryptAES(byte[] data, SecretKey key, byte[] iv) 
            throws GeneralSecurityException {
        if (data == null || key == null || iv == null) {
            throw new IllegalArgumentException("Data, key, and IV cannot be null");
        }
        
        Cipher cipher = Cipher.getInstance(AES_TRANSFORMATION);
        IvParameterSpec ivSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.ENCRYPT_MODE, key, ivSpec);
        
        byte[] encrypted = cipher.doFinal(data);
        
        return new EncryptionResult.Builder()
                .algorithm(EncryptionAlgorithm.AES)
                .data(encrypted)
                .metadata("iv", iv)
                .metadata("keySize", key.getEncoded().length * 8)
                .build();
    }
    
    /**
     * AES解密
     * 
     * @param encryptedData 加密数据
     * @param key 密钥
     * @param iv 初始化向量
     * @return 解密数据
     * @throws GeneralSecurityException 解密异常
     */
    public static byte[] decryptAES(byte[] encryptedData, SecretKey key, byte[] iv) 
            throws GeneralSecurityException {
        if (encryptedData == null || key == null || iv == null) {
            throw new IllegalArgumentException("Encrypted data, key, and IV cannot be null");
        }
        
        Cipher cipher = Cipher.getInstance(AES_TRANSFORMATION);
        IvParameterSpec ivSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.DECRYPT_MODE, key, ivSpec);
        
        return cipher.doFinal(encryptedData);
    }
    
    /**
     * DES加密
     * 
     * @param data 原始数据
     * @param key 密钥
     * @param iv 初始化向量
     * @return 加密结果
     * @throws GeneralSecurityException 加密异常
     */
    public static EncryptionResult encryptDES(byte[] data, SecretKey key, byte[] iv) 
            throws GeneralSecurityException {
        if (data == null || key == null || iv == null) {
            throw new IllegalArgumentException("Data, key, and IV cannot be null");
        }
        
        Cipher cipher = Cipher.getInstance(DES_TRANSFORMATION);
        IvParameterSpec ivSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.ENCRYPT_MODE, key, ivSpec);
        
        byte[] encrypted = cipher.doFinal(data);
        
        return new EncryptionResult.Builder()
                .algorithm(EncryptionAlgorithm.DES)
                .data(encrypted)
                .metadata("iv", iv)
                .metadata("keySize", key.getEncoded().length * 8)
                .build();
    }
    
    /**
     * DES解密
     * 
     * @param encryptedData 加密数据
     * @param key 密钥
     * @param iv 初始化向量
     * @return 解密数据
     * @throws GeneralSecurityException 解密异常
     */
    public static byte[] decryptDES(byte[] encryptedData, SecretKey key, byte[] iv) 
            throws GeneralSecurityException {
        if (encryptedData == null || key == null || iv == null) {
            throw new IllegalArgumentException("Encrypted data, key, and IV cannot be null");
        }
        
        Cipher cipher = Cipher.getInstance(DES_TRANSFORMATION);
        IvParameterSpec ivSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.DECRYPT_MODE, key, ivSpec);
        
        return cipher.doFinal(encryptedData);
    }
    
    /**
     * 3DES加密
     * 
     * @param data 原始数据
     * @param key 密钥
     * @param iv 初始化向量
     * @return 加密结果
     * @throws GeneralSecurityException 加密异常
     */
    public static EncryptionResult encryptTripleDES(byte[] data, SecretKey key, byte[] iv) 
            throws GeneralSecurityException {
        if (data == null || key == null || iv == null) {
            throw new IllegalArgumentException("Data, key, and IV cannot be null");
        }
        
        Cipher cipher = Cipher.getInstance(TRIPLE_DES_TRANSFORMATION);
        IvParameterSpec ivSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.ENCRYPT_MODE, key, ivSpec);
        
        byte[] encrypted = cipher.doFinal(data);
        
        return new EncryptionResult.Builder()
                .algorithm(EncryptionAlgorithm.TRIPLE_DES)
                .data(encrypted)
                .metadata("iv", iv)
                .metadata("keySize", key.getEncoded().length * 8)
                .build();
    }
    
    /**
     * 3DES解密
     * 
     * @param encryptedData 加密数据
     * @param key 密钥
     * @param iv 初始化向量
     * @return 解密数据
     * @throws GeneralSecurityException 解密异常
     */
    public static byte[] decryptTripleDES(byte[] encryptedData, SecretKey key, byte[] iv) 
            throws GeneralSecurityException {
        if (encryptedData == null || key == null || iv == null) {
            throw new IllegalArgumentException("Encrypted data, key, and IV cannot be null");
        }
        
        Cipher cipher = Cipher.getInstance(TRIPLE_DES_TRANSFORMATION);
        IvParameterSpec ivSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.DECRYPT_MODE, key, ivSpec);
        
        return cipher.doFinal(encryptedData);
    }
    
    /**
     * Blowfish加密
     * 
     * @param data 原始数据
     * @param key 密钥
     * @param iv 初始化向量
     * @return 加密结果
     * @throws GeneralSecurityException 加密异常
     */
    public static EncryptionResult encryptBlowfish(byte[] data, SecretKey key, byte[] iv) 
            throws GeneralSecurityException {
        if (data == null || key == null || iv == null) {
            throw new IllegalArgumentException("Data, key, and IV cannot be null");
        }
        
        Cipher cipher = Cipher.getInstance(BLOWFISH_TRANSFORMATION);
        IvParameterSpec ivSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.ENCRYPT_MODE, key, ivSpec);
        
        byte[] encrypted = cipher.doFinal(data);
        
        return new EncryptionResult.Builder()
                .algorithm(EncryptionAlgorithm.BLOWFISH)
                .data(encrypted)
                .metadata("iv", iv)
                .metadata("keySize", key.getEncoded().length * 8)
                .build();
    }
    
    /**
     * Blowfish解密
     * 
     * @param encryptedData 加密数据
     * @param key 密钥
     * @param iv 初始化向量
     * @return 解密数据
     * @throws GeneralSecurityException 解密异常
     */
    public static byte[] decryptBlowfish(byte[] encryptedData, SecretKey key, byte[] iv) 
            throws GeneralSecurityException {
        if (encryptedData == null || key == null || iv == null) {
            throw new IllegalArgumentException("Encrypted data, key, and IV cannot be null");
        }
        
        Cipher cipher = Cipher.getInstance(BLOWFISH_TRANSFORMATION);
        IvParameterSpec ivSpec = new IvParameterSpec(iv);
        cipher.init(Cipher.DECRYPT_MODE, key, ivSpec);
        
        return cipher.doFinal(encryptedData);
    }
    
    /**
     * RSA加密
     * 
     * @param data 原始数据
     * @param publicKey 公钥
     * @param useOAEP 是否使用OAEP填充
     * @return 加密结果
     * @throws GeneralSecurityException 加密异常
     */
    public static EncryptionResult encryptRSA(byte[] data, PublicKey publicKey, boolean useOAEP) 
            throws GeneralSecurityException {
        if (data == null || publicKey == null) {
            throw new IllegalArgumentException("Data and public key cannot be null");
        }
        
        String transformation = useOAEP ? RSA_OAEP_TRANSFORMATION : RSA_TRANSFORMATION;
        Cipher cipher = Cipher.getInstance(transformation);
        cipher.init(Cipher.ENCRYPT_MODE, publicKey);
        
        byte[] encrypted = cipher.doFinal(data);
        
        return new EncryptionResult.Builder()
                .algorithm(EncryptionAlgorithm.RSA)
                .data(encrypted)
                .metadata("encryptionType", EncryptionType.ASYMMETRIC)
                .metadata("keySize", publicKey.getEncoded().length * 8)
                .build();
    }
    
    /**
     * RSA解密
     * 
     * @param encryptedData 加密数据
     * @param privateKey 私钥
     * @param useOAEP 是否使用OAEP填充
     * @return 解密数据
     * @throws GeneralSecurityException 解密异常
     */
    public static byte[] decryptRSA(byte[] encryptedData, PrivateKey privateKey, boolean useOAEP) 
            throws GeneralSecurityException {
        if (encryptedData == null || privateKey == null) {
            throw new IllegalArgumentException("Encrypted data and private key cannot be null");
        }
        
        String transformation = useOAEP ? RSA_OAEP_TRANSFORMATION : RSA_TRANSFORMATION;
        Cipher cipher = Cipher.getInstance(transformation);
        cipher.init(Cipher.DECRYPT_MODE, privateKey);
        
        return cipher.doFinal(encryptedData);
    }
    
    /**
     * 字符串加密
     * 
     * @param text 原始文本
     * @param key 密钥
     * @param iv 初始化向量
     * @param algorithm 加密算法
     * @return Base64编码的加密字符串
     * @throws GeneralSecurityException 加密异常
     */
    public static String encryptString(String text, SecretKey key, byte[] iv, EncryptionAlgorithm algorithm) 
            throws GeneralSecurityException {
        if (text == null) {
            throw new IllegalArgumentException("Text cannot be null");
        }
        
        byte[] data = text.getBytes();
        EncryptionResult result;
        
        switch (algorithm) {
            case AES:
                result = encryptAES(data, key, iv);
                break;
            case DES:
                result = encryptDES(data, key, iv);
                break;
            case TRIPLE_DES:
                result = encryptTripleDES(data, key, iv);
                break;
            case BLOWFISH:
                result = encryptBlowfish(data, key, iv);
                break;
            default:
                throw new UnsupportedOperationException("Unsupported algorithm for string encryption: " + algorithm);
        }
        
        return Base64.getEncoder().encodeToString(result.getData());
    }
    
    /**
     * 字符串解密
     * 
     * @param encryptedText Base64编码的加密字符串
     * @param key 密钥
     * @param iv 初始化向量
     * @param algorithm 加密算法
     * @return 解密后的字符串
     * @throws GeneralSecurityException 解密异常
     */
    public static String decryptString(String encryptedText, SecretKey key, byte[] iv, EncryptionAlgorithm algorithm) 
            throws GeneralSecurityException {
        if (encryptedText == null) {
            throw new IllegalArgumentException("Encrypted text cannot be null");
        }
        
        byte[] encryptedData = Base64.getDecoder().decode(encryptedText);
        byte[] decryptedData;
        
        switch (algorithm) {
            case AES:
                decryptedData = decryptAES(encryptedData, key, iv);
                break;
            case DES:
                decryptedData = decryptDES(encryptedData, key, iv);
                break;
            case TRIPLE_DES:
                decryptedData = decryptTripleDES(encryptedData, key, iv);
                break;
            case BLOWFISH:
                decryptedData = decryptBlowfish(encryptedData, key, iv);
                break;
            default:
                throw new UnsupportedOperationException("Unsupported algorithm for string decryption: " + algorithm);
        }
        
        return new String(decryptedData);
    }
    
    /**
     * 计算数据的哈希值
     * 
     * @param data 原始数据
     * @param algorithm 哈希算法（MD5, SHA-1, SHA-256, SHA-512）
     * @return 哈希值
     * @throws NoSuchAlgorithmException 算法不支持
     */
    public static byte[] hash(byte[] data, String algorithm) throws NoSuchAlgorithmException {
        if (data == null) {
            throw new IllegalArgumentException("Data cannot be null");
        }
        
        MessageDigest digest = MessageDigest.getInstance(algorithm);
        return digest.digest(data);
    }
    
    /**
     * 计算字符串的哈希值
     * 
     * @param text 原始字符串
     * @param algorithm 哈希算法
     * @return Base64编码的哈希值
     * @throws NoSuchAlgorithmException 算法不支持
     */
    public static String hashString(String text, String algorithm) throws NoSuchAlgorithmException {
        if (text == null) {
            throw new IllegalArgumentException("Text cannot be null");
        }
        
        byte[] hash = hash(text.getBytes(), algorithm);
        return Base64.getEncoder().encodeToString(hash);
    }
    
    /**
     * 验证数据完整性
     * 
     * @param data 原始数据
     * @param expectedHash 期望的哈希值
     * @param algorithm 哈希算法
     * @return true如果验证通过
     * @throws NoSuchAlgorithmException 算法不支持
     */
    public static boolean verifyIntegrity(byte[] data, byte[] expectedHash, String algorithm) 
            throws NoSuchAlgorithmException {
        if (data == null || expectedHash == null) {
            return false;
        }
        
        byte[] actualHash = hash(data, algorithm);
        return Arrays.equals(actualHash, expectedHash);
    }
    
    /**
     * 生成消息认证码（HMAC）
     * 
     * @param data 原始数据
     * @param key 密钥
     * @param algorithm HMAC算法（HmacSHA256, HmacSHA512）
     * @return HMAC值
     * @throws NoSuchAlgorithmException 算法不支持
     * @throws InvalidKeyException 密钥无效
     */
    public static byte[] generateHMAC(byte[] data, SecretKey key, String algorithm) 
            throws NoSuchAlgorithmException, InvalidKeyException {
        if (data == null || key == null) {
            throw new IllegalArgumentException("Data and key cannot be null");
        }
        
        Mac mac = Mac.getInstance(algorithm);
        mac.init(key);
        return mac.doFinal(data);
    }
    
    /**
     * 验证消息认证码
     * 
     * @param data 原始数据
     * @param key 密钥
     * @param expectedHMAC 期望的HMAC值
     * @param algorithm HMAC算法
     * @return true如果验证通过
     * @throws NoSuchAlgorithmException 算法不支持
     * @throws InvalidKeyException 密钥无效
     */
    public static boolean verifyHMAC(byte[] data, SecretKey key, byte[] expectedHMAC, String algorithm) 
            throws NoSuchAlgorithmException, InvalidKeyException {
        if (data == null || key == null || expectedHMAC == null) {
            return false;
        }
        
        byte[] actualHMAC = generateHMAC(data, key, algorithm);
        return Arrays.equals(actualHMAC, expectedHMAC);
    }
    
    /**
     * 安全擦除字节数组
     * 
     * @param data 要擦除的数据
     */
    public static void secureErase(byte[] data) {
        if (data != null) {
            Arrays.fill(data, (byte) 0);
        }
    }
    
    /**
     * 安全擦除字符数组
     * 
     * @param data 要擦除的数据
     */
    public static void secureErase(char[] data) {
        if (data != null) {
            Arrays.fill(data, '\0');
        }
    }
    
    /**
     * 比较两个字节数组是否相等（防止时序攻击）
     * 
     * @param a 第一个字节数组
     * @param b 第二个字节数组
     * @return true如果相等
     */
    public static boolean constantTimeEquals(byte[] a, byte[] b) {
        if (a == null || b == null) {
            return a == b;
        }
        
        if (a.length != b.length) {
            return false;
        }
        
        int result = 0;
        for (int i = 0; i < a.length; i++) {
            result |= a[i] ^ b[i];
        }
        
        return result == 0;
    }
    
    /**
     * 获取算法的块大小
     * 
     * @param algorithm 加密算法
     * @return 块大小（字节）
     * @throws NoSuchAlgorithmException 算法不支持
     */
    public static int getBlockSize(EncryptionAlgorithm algorithm) throws NoSuchAlgorithmException {
        switch (algorithm) {
            case AES:
                return 16; // 128 bits
            case DES:
                return 8;  // 64 bits
            case TRIPLE_DES:
                return 8;  // 64 bits
            case BLOWFISH:
                return 8;  // 64 bits (variable, but 8 is common)
            case RSA:
                // RSA doesn't have a block size in the same sense
                // Returns the key size in bytes for RSA
                return -1;
            default:
                throw new NoSuchAlgorithmException("Unknown algorithm: " + algorithm);
        }
    }
    
    /**
     * 检查数据长度是否适合算法
     * 
     * @param data 数据
     * @param algorithm 加密算法
     * @return true如果适合
     * @throws NoSuchAlgorithmException 算法不支持
     */
    public static boolean isValidDataLength(byte[] data, EncryptionAlgorithm algorithm) 
            throws NoSuchAlgorithmException {
        if (data == null) {
            return false;
        }
        
        // 对于块加密算法，数据长度没有严格限制
        // 对于RSA，数据长度受密钥大小限制
        if (algorithm == EncryptionAlgorithm.RSA) {
            // RSA通常只能加密比密钥尺寸小的数据
            // 这里返回true，实际限制取决于密钥大小
            return true;
        }
        
        return true;
    }
}
